import os
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from dotenv import load_dotenv
import logging
import shutil

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 加载环境变量
load_dotenv()

# 定义模型名称
model_name = "D:/ideaSpace/MyPython/models/bge-small-zh-v1.5"

# 定义PDF路径
pdf_path = "D:/document/阿里巴巴Java开发手册-嵩山版.pdf"

# 加载PDF文件
try:
    loader = PyPDFLoader(pdf_path)
    documents = loader.load()
    logger.info("PDF文件加载成功")
except Exception as e:
    logger.error(f"PDF文件加载失败: {e}")
    raise

# 分割文档
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
texts = text_splitter.split_documents(documents)
logger.info("文档分割成功")

# 创建向量存储
try:
    embeddings = HuggingFaceEmbeddings(model_name=model_name)
    vectorstore = Chroma.from_documents(documents=texts, embedding=embeddings)
    logger.info("向量存储创建成功")
except Exception as e:
    logger.error(f"向量存储创建失败: {e}")
    raise

# 定义RAG系统
class ZhipuRAGSystem:
    def __init__(self, vectorstore):
        self.vectorstore = vectorstore
        self.chat_model = ChatZhipuAI(
            api_key=os.getenv("ZHIPUAI_API_KEY"),
            model_name="glm-4",
            temperature=0
        )
        self.prompt = ChatPromptTemplate.from_template(
            "根据以下上下文回答问题：\n\n{context}\n\n问题：{question}"
        )
        self.chain = (
                {"context": self.vectorstore.as_retriever(), "question": RunnablePassthrough()}
                | self.prompt
                | self.chat_model
                | StrOutputParser()
        )

    def query(self, question):
        return self.chain.invoke(question)

# 初始化RAG系统
try:
    rag = ZhipuRAGSystem(vectorstore)
    logger.info("RAG系统初始化成功")
except Exception as e:
    logger.error(f"RAG系统初始化失败: {e}")
    raise

# 交互式查询
print("\nRAG系统已就绪，输入'退出'结束")
while True:
    question = input("\n请输入问题: ").strip()
    if question.lower() in ['退出', 'exit', 'quit']:
        break

    if not question:
        print("问题不能为空")
        continue

    try:
        answer = rag.query(question)
        print(f"\n回答：{answer}")
    except Exception as e:
        logger.error(f"查询失败: {e}")
        print("查询失败，请稍后再试")
